clear
% addpath /Applications/Dynare/4.5.6/matlab
addpath C:\Users\schorf\Dropbox\Applications\Dynare\matlab
cd('./Auxiliary Functions');

%----------------------------------------------------------------
% Set parameters
%----------------------------------------------------------------
ParaSel = 1;

if ParaSel == 1
   setParameters;
   mom_1_1_shift = 0;
   mom_2_1_shift = 0;
   PhiepsC_shift = 0;
   PhiepsC_addsh = 0.01;
elseif ParaSel == 2
   setParameters;
   mom_1_1_shift = -0.15;
   mom_2_1_shift = -0.15;
   PhiepsC_shift = 0;
   PhiepsC_addsh = 0.01;
elseif ParaSel == 3
   setParameters;
   mom_1_1_shift = 0.7;
   mom_2_1_shift = 0.7;
   PhiepsC_shift = 0;
   PhiepsC_addsh = 0.01;     
end
   
%----------------------------------------------------------------
% Compute approximation tools
%----------------------------------------------------------------

% Grids
computeGrids;

% Polynomials over grids (only if using polynomials to approximate conditional expectation)
if splineOpt == 0
	computePolynomials;
end

%----------------------------------------------------------------
% Compute Steady State
%----------------------------------------------------------------

% Solve for steady state capital stock, distribution, and decision rules
coreSteadyState;

% Extract some important SS
K_star      = aggregateCapital; % in levels
mHat1_star  = mHat(1);
mHat2_star  = mHat(2);
mMean1_star = mMoments(1,1);
mMean2_star = mMoments(2,1); 

%----------------------------------------------------------------
% Save parameters in .mat files to import into Dynare 
%----------------------------------------------------------------

if splineOpt == 0	% if using polynomials to approximate individual decisions

	% Economic parameters
	save economicParameters.mat bbeta ssigma aaBar aalpha ddelta vEpsilonGrid aggEmployment ...
		mmu ttau rrhoTFP ssigmaTFP mEpsilonTransition
		
	% Approximation parameters
	save approximationParameters.mat nEpsilon nAssets nState assetsMin assetsMax nAssetsFine nStateFine nAssetsQuadrature nStateQuadrature ...
		nMeasure nMeasureCoefficients kRepSS maxIterations tolerance dampening
		
	% Grids
	save grids.mat vAssetsGridZeros vAssetsGrid mEpsilonGrid mAssetsGrid mEpsilonPrimeGrid vAssetsGridFine ...
		vAssetsGridFineZeros mEpsilonGridFine mAssetsGridFine mEpsilonPrimeGridFine vQuadratureWeights ...
		vAssetsGridQuadratureZeros vAssetsGridQuadrature mEpsilonGridQuadrature mAssetsGridQuadrature
		
	% Polynomials
	save polynomials.mat vAssetsPoly vAssetsPolySquared vAssetsPolyFine vAssetsPolyQuadrature vAssetsPolyBC
	
else	% if using splines to approximate individual decisions

	% Economic parameters
	save economicParameters.mat bbeta ssigma aaBar aalpha ddelta vEpsilonGrid aggEmployment ...
		mmu ttau rrhoTFP ssigmaTFP mEpsilonTransition
		
	% Approximation parameters
	save approximationParameters.mat nEpsilon nAssets nState assetsMin assetsMax nAssetsFine nStateFine nAssetsQuadrature nStateQuadrature ...
		nMeasure nMeasureCoefficients kRepSS maxIterations tolerance dampening
		
	% Grids
	save grids.mat vAssetsGrid mEpsilonGrid mAssetsGrid mEpsilonPrimeGrid vAssetsGridFine ...
		mEpsilonGridFine mAssetsGridFine mEpsilonPrimeGridFine vQuadratureWeights ...
		vAssetsGridQuadrature mEpsilonGridQuadrature mAssetsGridQuadrature
	
end

%----------------------------------------------------------------
% Run Dynare / Winberry Code
%----------------------------------------------------------------

if splineOpt == 0	% if using polynomials to approximate individual decisions

	dynare firstOrderDynamics_polynomials
	
else	% if using splines to approximate individual decisions

	dynare firstOrderDynamics_splines
	
end

cd('../')

save SimulDesign_data.mat mHat_1 mHat_2 moment_1_1 moment_2_1 moment_1_2 moment_2_2 moment_1_3 moment_2_3 ...
    measureCoefficient_1_1 measureCoefficient_2_1 measureCoefficient_1_2 measureCoefficient_2_2 measureCoefficient_1_3 measureCoefficient_2_3 ...
    r w aggregateTFP logAggregateOutput logAggregateInvestment logAggregateConsumption  

%%
%----------------------------------------------------------------
% Work with DYNARE output
%----------------------------------------------------------------

% extract coefficient matrices  (what's printed in the Dynare output "Policy and Transition Functions")
[ns, ne] = size(oo_.dr.ghx');
Phi1all   = zeros(ns, ne); % ns: number of state variables, ne: number of endogenous variables
Phiepsall = zeros(1, ne);
for ii = 1:ne
    Phi1all(:,ii)   = oo_.dr.ghx(oo_.dr.inv_order_var(ii),:)'; % 9x71 each column is an equation
    Phiepsall(:,ii) = oo_.dr.ghu(oo_.dr.inv_order_var(ii),:)'; % 1x71 each column is an equation
end
Phi1C   = Phi1all(:,[51 52 53 54 55 56 63 64 67])'; % each row is an equation
PhiepsC = Phiepsall(:,[51 52 53 54 55 56 63 64 67])'; % each row is an equation
% variables: mom_1_1, mom_2_1, mom_1_2, mom_2_2, mom_1_3, mom_2_3, mHat_1, mHat_2, aggTFP

Phi1mc   = Phi1all(:,[57 58 59 60 61 62])';
Phiepsmc = Phiepsall(:,[57 58 59 60 61 62])';
% variables: measureCoef_1_1, _2_1, ...

% extract steady states
ss_C  = oo_.dr.ys([51 52 53 54 55 56 63 64 67]);
ss_mc = oo_.dr.ys([57 58 59 60 61 62]);

%%
%----------------------------------------------------------------
% Adjust solution parameters to alter dynamics of the system
%----------------------------------------------------------------
Phi1C(1,3) = Phi1C(1,3) + mom_1_1_shift;
Phi1C(2,4) = Phi1C(2,4) + mom_2_1_shift;

if PhiepsC_shift == 1;
    PhiepsC_addsh_vec = PhiepsC_addsh*[1; 1; zeros(7,1)];
    PhiepsC = [PhiepsC PhiepsC_addsh_vec]
end


%%
%----------------------------------------------------------------
% Simulate aggregate data and cross-sectional distr. parameters
%----------------------------------------------------------------

load('./Auxiliary Functions/economicParameters.mat')

nvar_c  = 9;
nvar_r  = 2;
nvar_mc = 6;
T       = 280;
Tdrop   = 20;

% Create a selection matrix to extract TFP and K
M = zeros(nvar_c,nvar_r);
M(9,1) = 1;
M(7,2) = -(1-aggEmployment)*mMean1_star;
M(1,2) = (1-aggEmployment)*(1-mHat1_star);
M(8,2) = -aggEmployment*mMean2_star;
M(2,2) = aggEmployment*(1-mHat2_star);

% initalize output matrices and simulate
nsh    = size(PhiepsC,2);
ya_sim = zeros(nvar_c,T);
mc_sim = zeros(size(Phi1mc,1),T);

rng(1234);
for tt = 2:T
    sh_tt        = randn(nsh,1);
    ya_sim(:,tt) = Phi1C*ya_sim(:,tt-1) + PhiepsC*sh_tt;
    mc_sim(:,tt) = Phi1mc*ya_sim(:,tt-1) + Phiepsmc*sh_tt;
end
ya_sim      = ya_sim(:,Tdrop+1:end)';
aggvar_sim  = ya_sim*M;
mc_sim      = mc_sim(:,Tdrop+1:end)';
Tsim        = T - Tdrop;

sName     = ['Para' , num2str(ParaSel) ];
dataDir   = [pwd, '/', 'Data' ,'/', sName,'/'];
[~, ~, ~] = mkdir(dataDir);
csvwrite([dataDir, 'ya_sim.csv'], ya_sim);
csvwrite([dataDir, 'aggvar_sim.csv'], aggvar_sim);
csvwrite([dataDir, 'mc_sim.csv'], mc_sim);

% compute steady state K
ss_K = (1-aggEmployment)*(1-mHat1_star)*mMean1_star ...
       + aggEmployment*(1-mHat2_star)*mMean2_star;
csvwrite([dataDir, 'ss_K.csv'], ss_K);   

%%
%----------------------------------------------------------------
% Plot aggregate time series
%----------------------------------------------------------------

sName  = ['Para' , num2str(ParaSel) ];
figDir = [pwd, '/', 'Figures' ,'/', sName,'/'];
[~, ~, ~] = mkdir(figDir);

Tf = 1;
Tl = 160;

% Plot TFP
figure(1);clf;
set(figure(1),'PaperType','usletter','PaperOrientation','Landscape','PaperPosition',[0.1 0.1 11 8.5]);
plot(Tf:Tl,aggvar_sim(Tf:Tl,1)','Color','b','LineStyle','-','LineWidth',2)  
hold on
plot(Tf:Tl, zeros(Tl-Tf+1), 'Color', 'k', 'LineStyle', '-', 'LineWidth',2)
set(gca,'FontSize',50)
xlim([Tf Tl])

sFigName = [sName, '_TFPsim.pdf'];
saveas(figure(1), [figDir sFigName] );

% Plot K
figure(2);clf;
set(figure(2),'PaperType','usletter','PaperOrientation','Landscape','PaperPosition',[0.1 0.1 11 8.5]);
plot(Tf:Tl,aggvar_sim(Tf:Tl,2)+ss_K,'Color','b','LineStyle','-','LineWidth',4)  
hold on
plot(Tf:Tl, ss_K*ones(Tl-Tf+1), 'Color', 'k', 'LineStyle', '-', 'LineWidth',2)
hold on
plot(115*ones(Tl-Tf+1), linspace(2.5,3.5,Tl-Tf+1), 'Color', 'k', 'LineStyle', '--', 'LineWidth',2)
hold on
plot(155*ones(Tl-Tf+1), linspace(2.5,3.5,Tl-Tf+1), 'Color', 'k', 'LineStyle', '--', 'LineWidth',2)
set(gca,'FontSize',50)
xlim([Tf Tl])
ylim([2.5 3.5])

sFigName = [sName, '_Ksim.pdf'];
saveas(figure(2), [figDir sFigName] );


%%
%----------------------------------------------------------------
% Simulate cross-sectional data 
%----------------------------------------------------------------
Nsim    = 10000;
mom_sim_ss = ya_sim(:,1:6) + ss_C(1:6,1)';
mc_sim_ss  = mc_sim + ss_mc';

% asset holding grid
a_min  = 0;
a_max  = 8;
a_n    = 100;
a_grid = linspace(a_min,a_max, a_n)';

% initialize output matrices and simulate
a_cross_sim    = zeros(Nsim,Tsim);
empl_cross_sim = zeros(Nsim,Tsim);

rng(1234);
for tt = 1:Tsim
    
    mom_1_1_tt = mom_sim_ss(tt,1);
    mom_2_1_tt = mom_sim_ss(tt,2);
    mom_1_2_tt = mom_sim_ss(tt,3);
    mom_2_2_tt = mom_sim_ss(tt,4);
    mom_1_3_tt = mom_sim_ss(tt,5);
    mom_2_3_tt = mom_sim_ss(tt,6);
    
    mc_1_1_tt = mc_sim_ss(tt,1);
    mc_2_1_tt = mc_sim_ss(tt,2);
    mc_1_2_tt = mc_sim_ss(tt,3);
    mc_2_2_tt = mc_sim_ss(tt,4);
    mc_1_3_tt = mc_sim_ss(tt,5);
    mc_2_3_tt = mc_sim_ss(tt,6);

    unemployed_g = @(a) exp(mc_1_1_tt * (a - mom_1_1_tt) ...
        + mc_1_2_tt * ((a - mom_1_1_tt).^2 - mom_1_2_tt) ...
        + mc_1_3_tt * ((a - mom_1_1_tt).^3 - mom_1_3_tt));
    employed_g   = @(a) exp(mc_2_1_tt * (a - mom_2_1_tt) ...
        + mc_2_2_tt * ((a - mom_2_1_tt).^2 - mom_2_2_tt) ...
        + mc_2_3_tt * ((a - mom_2_1_tt).^3 - mom_2_3_tt));
    
    unemployed_norm = integral(unemployed_g,a_min,a_max);
    employed_norm   = integral(employed_g,a_min, a_max);
    
    unemployed_pdf = @(a) unemployed_g(a)/unemployed_norm;
    employed_pdf   = @(a) employed_g(a)/employed_norm;
    
    unemployed_cdf = @(a) integral(@(x) unemployed_pdf(x),a_min,a);
    employed_cdf   = @(a) integral(@(x) employed_pdf(x),a_min,a);
    
    empl_cross_sim(:,tt) = (rand(Nsim,1) < aggEmployment) ;
    uu = rand(Nsim);
    
    parfor ii = 1:Nsim
        if empl_cross_sim(ii,tt) == 1
            a_cross_sim(ii,tt) = fzero(@(a) employed_cdf(a)-uu(ii), mom_2_1_tt);
        else
            a_cross_sim(ii,tt) = fzero(@(a) unemployed_cdf(a)-uu(ii), mom_2_1_tt);
        end   
    end
    
    sprintf('t= %d',tt)

end

sName     = ['Para' , num2str(ParaSel) ];
dataDir   = [pwd, '/', 'Data' ,'/', sName,'/'];
[~, ~, ~] = mkdir(dataDir);
csvwrite([dataDir, 'a_cross_sim.csv'], a_cross_sim');
csvwrite([dataDir, 'empl_cross_sim.csv'], empl_cross_sim');

%%
%----------------------------------------------------------------
% Plot cross-sectional distribution 
%----------------------------------------------------------------

% we need to add steady states and we will start in the steady state
mom_sim_ss = ya_sim(:,1:6) + ss_C(1:6,1)';
mc_sim_ss  = mc_sim + ss_mc';

% asset holding grid
a_min  = 0;
a_max  = 8;
a_n    = 100;
a_grid = linspace(a_min,a_max, a_n)';

% choose time period
tt = 45;

mom_1_1_tt = mom_sim_ss(tt,1);
mom_2_1_tt = mom_sim_ss(tt,2);
mom_1_2_tt = mom_sim_ss(tt,3);
mom_2_2_tt = mom_sim_ss(tt,4);
mom_1_3_tt = mom_sim_ss(tt,5);
mom_2_3_tt = mom_sim_ss(tt,6);

mc_1_1_tt = mc_sim_ss(tt,1);
mc_2_1_tt = mc_sim_ss(tt,2);
mc_1_2_tt = mc_sim_ss(tt,3);
mc_2_2_tt = mc_sim_ss(tt,4);
mc_1_3_tt = mc_sim_ss(tt,5);
mc_2_3_tt = mc_sim_ss(tt,6);

unemployed_g = @(a) exp(mc_1_1_tt * (a - mom_1_1_tt) ...
    + mc_1_2_tt * ((a - mom_1_1_tt).^2 - mom_1_2_tt) ...
    + mc_1_3_tt * ((a - mom_1_1_tt).^3 - mom_1_3_tt));
employed_g   = @(a) exp(mc_2_1_tt * (a - mom_2_1_tt) ...
    + mc_2_2_tt * ((a - mom_2_1_tt).^2 - mom_2_2_tt) ...
    + mc_2_3_tt * ((a - mom_2_1_tt).^3 - mom_2_3_tt));

unemployed_norm = integral(unemployed_g,a_min,a_max);
employed_norm   = integral(employed_g,a_min, a_max);

unemployed_pdf = @(a) unemployed_g(a)/unemployed_norm;
employed_pdf   = @(a) employed_g(a)/employed_norm;

a_cross_pdf_unempl = unemployed_pdf(a_grid);
a_cross_pdf_empl   = employed_pdf(a_grid);
a_cross_pdf_all    = aggEmployment*a_cross_pdf_empl ...
                   +(1-aggEmployment)*a_cross_pdf_unempl;

% Plot densities of a 
sName   = ['Para' , num2str(ParaSel) ];
figDir  = [pwd, '/', 'Figures' ,'/', sName,'/'];
[~, ~, ~] = mkdir(figDir);

figure(1);clf;
set(figure(1),'PaperType','usletter','PaperOrientation','Landscape','PaperPosition',[0.1 0.1 11 8.5]);
plot(a_grid,a_cross_pdf_unempl,'Color','b','LineStyle','--','LineWidth',2)  
hold on
plot(a_grid,a_cross_pdf_empl,'Color','r','LineStyle',':','LineWidth',2)  
hold on
plot(a_grid,a_cross_pdf_all, 'Color', 'g', 'LineStyle', '-', 'LineWidth',2)
set(gca,'FontSize',50)
xlim([a_min a_max])

sFigName = [sName, '_a_cross_t_',num2str(tt),'.pdf'];
saveas(figure(1), [figDir sFigName] );

% Plot densities of a/K
ss_K = (1-aggEmployment)*(1-mHat1_star)*mMean1_star ...
       + aggEmployment*(1-mHat2_star)*mMean2_star;

K_tt = aggvar_sim(tt,2)+ss_K;
aK_cross_pdf_unempl = a_cross_pdf_unempl*K_tt;
aK_cross_pdf_empl   = a_cross_pdf_empl*K_tt;
aK_cross_pdf_all    = a_cross_pdf_all*K_tt;

aK_grid = a_grid/K_tt;

sName     = ['Para' , num2str(ParaSel) ];
dataDir   = [pwd, '/', 'Data' ,'/', sName,'/'];
[~, ~, ~] = mkdir(dataDir);
csvwrite([dataDir, 'aK_cross_pdf_empl_tt', num2str(tt), '.csv'], [aK_grid aK_cross_pdf_empl]);

figure(2);clf;
set(figure(2),'PaperType','usletter','PaperOrientation','Landscape','PaperPosition',[0.1 0.1 11 8.5]);
plot(aK_grid,aK_cross_pdf_unempl,'Color','r','LineStyle','--','LineWidth',4)  
hold on
plot(aK_grid,aK_cross_pdf_empl,'Color','b','LineStyle','-','LineWidth',4)  
hold on
plot(aK_grid,aK_cross_pdf_all, 'Color', 'k', 'LineStyle', '-', 'LineWidth',2)
hold on
plot(ones(a_n),linspace(0,1.2,a_n), 'Color', 'k', 'LineStyle', ':', 'LineWidth',2)
set(gca,'FontSize',50)
% xlim([a_min/K_tt a_max/K_tt])
xlim([0 2])
ylim([0 1.2])

sFigName = [sName, '_aK_cross_t_',num2str(tt),'.pdf'];
saveas(figure(2), [figDir sFigName] );

% Overlay histogram and density of a
figure(3);clf;
set(figure(3),'PaperType','usletter','PaperOrientation','Landscape','PaperPosition',[0.1 0.1 11 8.5]);
plot(a_grid,a_cross_pdf_all,'Color','b','LineStyle','--','LineWidth',2)  
hold on
histogram(a_cross_sim(:,tt), 40, 'Normalization', 'pdf','FaceColor',1/255*[200,200,200],'EdgeColor','none')
set(gca,'FontSize',50)
xlim([a_min a_max])

sFigName = [sName, '_a_cross_pdf_hist_t_',num2str(tt),'.pdf'];
saveas(figure(3), [figDir sFigName] );

% Overlay histogram and density of a/K
figure(4);clf;
set(figure(4),'PaperType','usletter','PaperOrientation','Landscape','PaperPosition',[0.1 0.1 11 8.5]);
plot(aK_grid,aK_cross_pdf_all,'Color','b','LineStyle','--','LineWidth',2)  
hold on
histogram(a_cross_sim(:,tt)/K_tt, 40, 'Normalization', 'pdf','FaceColor',1/255*[200,200,200],'EdgeColor','none')
set(gca,'FontSize',50)
xlim([a_min/K_tt a_max/K_tt])

sFigName = [sName, '_aK_cross_pdf_hist_t_',num2str(tt),'.pdf'];
saveas(figure(4), [figDir sFigName] );


